import numpy as np
import matplotlib.pyplot as plt

# 参数设置
a = 1.0       # x方向长度
b = 1.0       # y方向长度
u0 = 0.0      # 冷却介质温度
U = 100.0     # 高温边界温度
N = 50        # 级数项数 (k=0 到 k=N-1)
nx, ny = 200, 200  # 网格分辨率

# 创建网格
x = np.linspace(0, a, nx)
y = np.linspace(0, b, ny)
X, Y = np.meshgrid(x, y)

# 初始化温度场
u = np.full_like(X, u0)

# 计算级数项
for k in range(N):
    m = 2 * k + 1
    alpha = m * np.pi / a
    # 双曲正弦比值（避免大数溢出）
    sinh_ratio = np.sinh(alpha * Y) / np.sinh(alpha * b)
    # 更新温度场
    u += (1 / m) * sinh_ratio * np.sin(alpha * X)

# 乘以系数
u = u0 + (a * (U - u0) / np.pi) * u

# 绘图
plt.figure(figsize=(10, 8), dpi=300)
im = plt.pcolormesh(X, Y, u, cmap='jet', shading='auto')
plt.colorbar(im, label='Temperature (°C)')
plt.xlabel('x')
plt.ylabel('y')
plt.title('Steady-State Temperature Distribution in Heat Sink')
plt.axis('equal')

# 保存图像
plt.savefig('热平衡8-1例题3.png', dpi=300, bbox_inches='tight')
plt.close()